package com.itcast.peft.controller;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.Random;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.ResponseBody;
import org.springframework.web.bind.annotation.RestController;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.itcast.peft.entity.LineMessage;
import com.itcast.peft.entity.RoleMessage;

@RestController
@RequestMapping("/sft")
public class SFTController {

    /**
     * objectMapper用于序列化
     */
    private ObjectMapper objectMapper = new ObjectMapper();
    
    /**
     * 模型对象
     */
    @Autowired
    private ChatModel chatModel;

    /**
     * 常见大型动物数组，使用数组是为了方便随机筛选
     */
    private String[] largeAnimals = new String[] {"狮子","鲸鱼","犀牛","河马","长颈鹿","熊","鳄鱼","斑马","骆驼","水牛"};
    /**
     * jsonl数据生成
     * @return
     */
    @GetMapping("/generateJsonLine")
    @ResponseBody
    public ResponseEntity<byte[]> generateJsonLine() {
        try (ByteArrayOutputStream bios = new ByteArrayOutputStream()) {
            // 初始化chatMemory
            var mem = MessageWindowChatMemory.builder().build();
            // 初始化cli
            var cli = ChatClient.builder(chatModel)
                    .defaultAdvisors(MessageChatMemoryAdvisor.builder(mem).build())
                    // 注意，我们需要简洁的回答适配小模型训练，但可能会影响教师大模型的发挥，因此我们使用Prompt进行预设
                    .defaultSystem("你需要回答关于动物放进冰箱的脑筋急转弯，初始时冰箱是空的，你要注意如果冰箱里已有其他动物需要先拿出来，你的回答非常简洁，不解释，只输出结果")
                    .build();
            // 本次我们生成10条数据
            for (int i = 0; i < 10; i++) {
                ask(cli);
                var messages = mem.get(ChatMemory.DEFAULT_CONVERSATION_ID).stream()
                        .filter(m -> m.getMessageType() != MessageType.SYSTEM)
                        .map(RoleMessage::fromNormalMessage)
                        .toList();
                String jsonLine = objectMapper.writeValueAsString(new LineMessage(messages)) + '\n';
                bios.writeBytes(jsonLine.getBytes());
                mem.clear(ChatMemory.DEFAULT_CONVERSATION_ID);
            }
            // 输出数据为下载文件的形式
            return ResponseEntity.ok()
                    .contentType(MediaType.APPLICATION_JSON)
                    .header("Content-disposition", "attachment; filename=data.jsonl")
                    .body(bios.toByteArray());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }
    
    private void ask(ChatClient cli) {
        // 选择随机动物
        Random r = new Random();
        int idxA = r.nextInt(largeAnimals.length);
        int idxB = r.nextInt(largeAnimals.length);
        while (idxA == idxB) {
            idxB = r.nextInt(largeAnimals.length);
        }
        // 进行教师模型提问
        cli.prompt().user(String.format("把%s放到冰箱分几步？", largeAnimals[idxA])).call().content();
        cli.prompt().user(String.format("把%s放到冰箱分几步？", largeAnimals[idxB])).call().content();
        cli.prompt().user("动物园开会谁没来？").call().content();
    }
}
